
import sys
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QTextEdit, QFileDialog, QLabel, QGridLayout
)
from PyQt6.QtCore import Qt
import pyqtgraph as pg
from obspy import read
import numpy as np

# 设置pyqtgraph的背景和前景颜色，使其更美观
pg.setConfigOption('background', 'w')
pg.setConfigOption('foreground', 'k')

# 自定义 ViewBox，禁用鼠标滚轮缩放（避免整体缩放）
class NoWheelViewBox(pg.ViewBox):
    def wheelEvent(self, ev):
        # 禁用滚轮缩放
        ev.accept()
        return

class SacViewerPyQtGraph(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("高性能SAC波形查看器 (PyQtGraph)")
        self.setGeometry(100, 100, 1200, 700)

        self.trace = None

        # --- 创建主布局 ---
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        grid_layout = QGridLayout(central_widget)

        # --- 顶部控件 ---
        controls_widget = QWidget()
        controls_layout = QHBoxLayout(controls_widget)
        
        self.load_button = QPushButton("加载SAC文件")
        self.reset_view_button = QPushButton("重置视图(全部)")
        self.filepath_label = QLabel("未加载文件。")
        self.filepath_label.setStyleSheet("color: gray;")

        # 轴向缩放控件（独立的 X/Y 缩放）
        self.zoom_in_x_btn = QPushButton("时间轴放大")
        self.zoom_out_x_btn = QPushButton("时间轴缩小")
        self.reset_x_btn = QPushButton("时间轴重置")

        self.zoom_in_y_btn = QPushButton("数值轴放大")
        self.zoom_out_y_btn = QPushButton("数值轴缩小")
        self.reset_y_btn = QPushButton("数值轴重置")

        controls_layout.addWidget(self.load_button)
        controls_layout.addWidget(self.reset_view_button)
        controls_layout.addSpacing(20)
        controls_layout.addWidget(self.zoom_in_x_btn)
        controls_layout.addWidget(self.zoom_out_x_btn)
        controls_layout.addWidget(self.reset_x_btn)
        controls_layout.addSpacing(10)
        controls_layout.addWidget(self.zoom_in_y_btn)
        controls_layout.addWidget(self.zoom_out_y_btn)
        controls_layout.addWidget(self.reset_y_btn)
        controls_layout.addStretch(1)

        # --- 绘图区域 ---
        axis_items = {'bottom': pg.DateAxisItem()}
        self.plot_widget = pg.PlotWidget(viewBox=NoWheelViewBox(), axisItems=axis_items)
        self.plot_widget.getPlotItem().showGrid(x=True, y=True, alpha=0.5)
        self.plot_widget.setLabel('left', "振幅")
        self.plot_widget.setLabel('bottom', "时间")

        # 仅允许鼠标拖动平移，不启用滚轮缩放（滚轮已禁用）
        self.plot_widget.getViewBox().setMouseEnabled(x=True, y=True)

        # --- SAC头信息显示 ---
        self.header_text = QTextEdit()
        self.header_text.setReadOnly(True)
        self.header_text.setLineWrapMode(QTextEdit.LineWrapMode.NoWrap)
        
        header_label = QLabel("SAC 头信息:")

        # --- 布局安排 ---
        grid_layout.addWidget(controls_widget, 0, 0, 1, 2)
        grid_layout.addWidget(self.filepath_label, 1, 0, 1, 2)
        grid_layout.addWidget(self.plot_widget, 2, 0)
        
        header_layout = QVBoxLayout()
        header_layout.addWidget(header_label)
        header_layout.addWidget(self.header_text)
        grid_layout.addLayout(header_layout, 2, 1)

        grid_layout.setColumnStretch(0, 3)  # 绘图区域
        grid_layout.setColumnStretch(1, 1)  # 头信息区域

        # --- 信号连接 ---
        self.load_button.clicked.connect(self.load_sac_file)
        self.reset_view_button.clicked.connect(self.reset_view)

        self.zoom_in_x_btn.clicked.connect(lambda: self.zoom_x(factor=1.5))
        self.zoom_out_x_btn.clicked.connect(lambda: self.zoom_x(factor=1/1.5))
        self.reset_x_btn.clicked.connect(self.reset_x)

        self.zoom_in_y_btn.clicked.connect(lambda: self.zoom_y(factor=1.5))
        self.zoom_out_y_btn.clicked.connect(lambda: self.zoom_y(factor=1/1.5))
        self.reset_y_btn.clicked.connect(self.reset_y)

    def load_sac_file(self):
        filepath, _ = QFileDialog.getOpenFileName(
            self,
            "选择一个SAC文件",
            "",
            "SAC Files (*.sac);;All Files (*)"
        )
        if not filepath:
            return

        try:
            stream = read(filepath)
            if not stream:
                self.filepath_label.setText(f"错误: 无法读取文件 {filepath}")
                return
            
            self.trace = stream[0]
            self.filepath_label.setText(filepath)
            
            self.update_plot()
            self.display_header()
            self.reset_view()

        except Exception as e:
            self.filepath_label.setText(f"错误: {e}")

    def update_plot(self):
        if self.trace is None:
            return

        data = self.trace.data
        timestamps = self.trace.times("timestamp")

        self.plot_widget.clear()
        self.curve = self.plot_widget.plot(timestamps, data, pen=pg.mkPen('k', width=1), antialias=True)
        
    def display_header(self):
        if self.trace is None:
            return
        
        header_info = ""
        if hasattr(self.trace.stats, 'sac'):
            for key, value in self.trace.stats.sac.items():
                header_info += f"{key:<10}: {value}\n"
        else:
            header_info = str(self.trace.stats)
            
        self.header_text.setText(header_info)

    # --- 视图控制 ---
    def reset_view(self):
        # 重置全部范围
        self.plot_widget.enableAutoRange(axis=pg.ViewBox.XYAxes, enable=True)
        self.plot_widget.autoRange()

    def reset_x(self):
        # 仅重置 X 轴范围
        vb = self.plot_widget.getViewBox()
        vb.enableAutoRange(axis=pg.ViewBox.XAxis, enable=True)
        vb.updateAutoRange()

    def reset_y(self):
        # 仅重置 Y 轴范围
        vb = self.plot_widget.getViewBox()
        vb.enableAutoRange(axis=pg.ViewBox.YAxis, enable=True)
        vb.updateAutoRange()

    def zoom_x(self, factor: float):
        # 围绕当前可视窗口中心进行 X 轴缩放
        vb = self.plot_widget.getViewBox()
        (x0, x1), (y0, y1) = vb.viewRange()
        x_center = 0.5 * (x0 + x1)
        x_len = (x1 - x0)
        new_len = x_len / factor if factor > 0 else x_len
        vb.setXRange(x_center - new_len/2, x_center + new_len/2, padding=0)

    def zoom_y(self, factor: float):
        # 围绕当前可视窗口中心进行 Y 轴缩放
        vb = self.plot_widget.getViewBox()
        (x0, x1), (y0, y1) = vb.viewRange()
        y_center = 0.5 * (y0 + y1)
        y_len = (y1 - y0)
        new_len = y_len / factor if factor > 0 else y_len
        vb.setYRange(y_center - new_len/2, y_center + new_len/2, padding=0)


def main():
    app = QApplication(sys.argv)
    viewer = SacViewerPyQtGraph()
    viewer.show()
    sys.exit(app.exec())

if __name__ == '__main__':
    main()
